from baselines.demo.HMM import HMM, TextSequence, viterbi
from typing import *
import pickle
import numpy as np
import re
import os

basedir = os.path.dirname(os.path.abspath(__file__))

if __name__ == '__main__':
    filename = os.path.join(basedir, "./pos.txt")
    dictname = os.path.join(basedir, "./dict.txt")
    modelname = os.path.join(basedir, "./model.hmm")
    print("数据读取中...")
    data = np.loadtxt(filename, dtype="str", delimiter="__label__", encoding="utf-8")
    #print(data.shape) # (45957,)

    def split_sentence_pos_row(row: str)->List[str]:
        """把一行标注数据转换成干净的列表格式"""
        return list(filter(lambda x: x!="", re.split(r"\s+", row)[1:]))

    """
        统计tag
        B-tag
        I
        O
    """
    inner_tag = "I"
    outer_Tag = "O"
    tags = {inner_tag}
    tag2idx = dict()
    idx2tag = []
    data2 = []
    for i in range(len(data)):
        sentence_frag = split_sentence_pos_row(data[i])
        sentence_char = []
        sentence_tags = []
        for frag in sentence_frag:
            word_tag = re.split(r"/", frag)
            if(len(word_tag) != 2 or word_tag[0] == '' or word_tag[1] == ''): #如果格式不符合  word/tag
                continue
            word, tag = word_tag
            b_tag = "B-" + tag
            if(not b_tag in tags):
                tags.add(b_tag)

            #把词分成字重新构造句子
            for i in range(len(word)):
                sentence_char.append(word[i])
                sentence_tags.append(b_tag if i == 0 else inner_tag)
        data2.append((sentence_char, sentence_tags))

    for tag in tags:
        tag2idx[tag] = len(idx2tag)
        idx2tag.append(tag)


    """
        使用BIO标注法对数据标注
        构造训练集
    """
    x = []
    y = []

    for chars, tags in data2:
        x.append(chars)
        y.append([ tag2idx[tag] for tag in tags])

    print("数据读取完毕")
    print("开始训练...")

    ### 测试算法
    hmm = TextSequence(x, y, len(idx2tag))
    print("训练完毕")
    print("正在保存字典...")
    np.savetxt(dictname,  np.array(idx2tag, dtype="str"), fmt="%s")
    print("字典保存完毕")
    print("正在保存模型...")
    with open(file=modelname, mode="bw+") as file:
        pickle.dump(hmm, file)
    print("模型保存成功")
    print("正在重新加载模型...")
    with open(file=modelname, mode="br") as file:
        hmm2 = pickle.load(file)
    print("模型加载成功")
    print("开始预测...")
    states = viterbi(hmm2, hmm2.getIndex(list("中华人民共和国今天成立了！")))
    print("预测结果为：")
    print([idx2tag[idx] for idx in states])